import pandas as pd
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
from scipy.cluster.hierarchy import linkage, fcluster
from sklearn.preprocessing import StandardScaler
from scipy.spatial.distance import pdist
from utils.utils import *
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False  # 正确显示负号

def data_clean(df):
    print("原始数据形状：", df.shape)
    # Step 2: 数据清洗（关键！）
    df.replace([np.inf, -np.inf], np.nan, inplace=True)
    missing_before = df.isnull().sum().sum()
    print(f"NaN/Inf 替换后，缺失值数量：{missing_before}")

    # 用均值填充 NaN
    df.fillna(df.mean(), inplace=True)
    missing_after = df.isnull().sum().sum()
    print(f"填充后剩余缺失值数量（应为0）：{missing_after}")
    if not np.isfinite(df.values).all():
        raise ValueError("数据中仍然存在 NaN 或 Inf，请检查原始数据！")
    return df
def cluster(df,type):
    df = data_clean(df)
    labels = df[target_column]
    data = df.drop(columns=[target_column,'Patient_ID'])  # 去除标签列，仅保留特征

    # 3. 标准化特征
    scaler = StandardScaler()
    X_scaled = scaler.fit_transform(data)

    Z = linkage(X_scaled, method='ward')
    clusters = fcluster(Z, t=2, criterion='maxclust')  # 分成2类
    cluster_labels = pd.Series(clusters, name="Cluster")

    # 5. 计算 Cluster 1 中各特征的平均值并排序
    X_df = pd.DataFrame(X_scaled, columns=data.columns, index=cluster_labels.index)
    cluster1_mean = X_df[cluster_labels == 1].mean(axis=0)
    sorted_feature_names = cluster1_mean.sort_values(ascending=False).index  # 排序后的特征名

    # 6. 热图绘制
    sns.set(style="white")
    d = pdist(X_scaled)
    linkage_matrix = linkage(d, method="ward")

    # 将 cluster label 转换为颜色条（1类=浅蓝，2类=深蓝）
    col_colors = cluster_labels.map({1: "#AEDFF7", 2: "#023E8A"}).to_numpy()

    # 对 X_scaled.T 进行排序
    X_sorted = pd.DataFrame(X_scaled, columns=data.columns).T.loc[sorted_feature_names].to_numpy()

    # 绘制聚类热图
    g = sns.clustermap(
        X_sorted,
        col_linkage=linkage_matrix,
        row_cluster=False,
        col_cluster=True,
        col_colors=col_colors,
        cmap="Blues",
        cbar_pos=None,
        vmin=0,
        vmax=2,
        xticklabels=False,
        yticklabels=False,
        figsize=(10, 10)
    )
    g.ax_row_dendrogram.set_visible(False)  # 彻底隐藏行树状图

    # 调整整个热图布局，让图更紧凑
    g.fig.subplots_adjust(left=0, right=0.95, top=0.95, bottom=0.05)
    plt.suptitle("A. Clustered matrix for multi-sequence profiles of ALN", y=1.02)
    save_path = opj(result_path,'cluster',type)
    md(save_path)
    plt.savefig(opj(save_path,'cluster_heatmap.pdf'))

    # 6. 合并标签和聚类结果
    result_df = pd.DataFrame({
        "Cluster": cluster_labels.map({1: "Cluster A", 2: "Cluster B"}),
        "ALN_metastasis": labels.map({1: "Yes", 0: "No", "Yes": "Yes", "No": "No"})  # 兼容各种格式
    })

    # 7. 画柱状图
    # 第一步：先计算数量和百分比
    count_df = (
        result_df
        .groupby(["Cluster", "ALN_metastasis"])
        .size()
        .reset_index(name='count')  # 变成 DataFrame
        .pivot_table(index="Cluster", columns="ALN_metastasis", values="count", fill_value=0)
    )

    percent_df = count_df.apply(lambda x: x / x.sum(), axis=1)  # 按行归一化

    # 第二步：画堆叠柱状图
    ax = percent_df[["No", "Yes"]].plot(kind="bar", stacked=True, color=["#003f5c", "#ffa600"])
    plt.ylabel("Percent weight")
    plt.title("B. Bar chart for cluster of ALN")
    plt.legend(title="ALN metastasis")
    plt.xticks(rotation=0)
    plt.tight_layout()

    # 第三步：添加百分比 + 人数
    for i, col in enumerate(["No", "Yes"]):
        for j, cluster in enumerate(percent_df.index):
            percent = percent_df.loc[cluster, col]
            count = count_df.loc[cluster, col]
            if percent > 0:
                ax.annotate(f'{percent:.1%}\n({count})',
                            (j, percent_df.loc[cluster, ["No"]].iloc[0] + percent / 2) if col == "Yes"
                            else (j, percent / 2),
                            ha='center', va='center', color='white' if col == "No" else 'black', fontsize=10)

    plt.savefig(opj(save_path,'cluster_distribution.pdf'))

    # 8. 统计显著性检验：卡方检验
    from scipy.stats import chi2_contingency

    # 构造列联表
    contingency_table = pd.crosstab(result_df["Cluster"], result_df["ALN_metastasis"])
    print("Contingency Table:")
    print(contingency_table)

    # 如果是2x2，且样本量较小，可使用 Fisher 精确检验；否则用卡方
    chi2, chi2_p, dof, expected = chi2_contingency(contingency_table)
    print(f"Chi-square Test P-value: {chi2_p:.4f}")
    if chi2_p < 0.05:
        print("✅ 显著相关：聚类与转移状态具有统计学相关性")
    else:
        print("❌ 不显著：聚类与转移状态无统计学相关性")

for t in data_to_solve_list:
    df = pd.read_csv(opj(merge_data_path,t+'.csv'))
    cluster(df,t)